import os
import glob
import argparse
from pathlib import Path
import gmalg

def sm9_dec(mpk_e_file, msk_e_file, sk_e_file, uid_file):
    hid_e = b"\x03"
    mpk_e_temp = read_binary_file_(mpk_e_file)
    mpk_e = b"\x04" + mpk_e_temp
    msk_e = read_binary_file_(msk_e_file)
    sk_e_temp = read_binary_file_(sk_e_file)
    sk_e = b"\x04" + sk_e_temp
    uid  = read_binary_file_(uid_file)
    #kgc = gmalg.SM9KGC(hid_e=hid_e, msk_e=msk_e, mpk_e=mpk_e)
    #uid = bytes.fromhex("8DFF2AB20AC81EA8A5BEA2EAC3B2A7F4")
    #sk_e = kgc.generate_sk_encrypt(uid)
    sm9 = gmalg.SM9(hid_e=hid_e, mpk_e=mpk_e, sk_e=sk_e, uid=uid)
    plain = b"Chinese IBE standard"
    cipher = sm9.encrypt(plain, uid)  # encrypt data to self
    assert sm9.decrypt(cipher) == plain, "验证失败" 
    #print("id", uid.hex())
    #print("msk_e", msk_e.hex())
    #print("mpk_e", mpk_e.hex())
    #print("sk_e", sk_e.hex())
    #print("plain", plain.hex())
    #print("cipher", cipher.hex())
    #print("replain", sm9.decrypt(cipher).hex())

def read_binary_file_(file_path):
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        return binary_data


def read_binary_file(file_path):
    """读取二进制文件内容并转换为字符串"""
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        # 尝试解码为字符串，如果失败则用十六进制表示
        return binary_data.hex().upper()

def main():
    parser = argparse.ArgumentParser(description='遍历目录中的二进制文件')
    parser.add_argument('directory', help='目标目录路径')
    parser.add_argument('mode', choices=['ok', 'speed'], default='ok', help='运行模式')
    args = parser.parse_args()
    directory = args.directory
    mode = args.mode

    if not os.path.isdir(directory):
        print(f"错误：目录不存在 - {directory}")
        return
    
    id_files = sorted(glob.glob(os.path.join(directory, 'id_*')))
    count = len(id_files)
    print(f"count={count}")
    # 创建输出内容
    merged_content = ""
    
    # 遍历0-9，合并对应的公钥和私钥文件
    for i in range(count):
        # 读取文件内容
        id_file = os.path.join(directory, f'id_{i}')
        Ke_file = os.path.join(directory, f'ke_{i}')
        Ppub_e_file = os.path.join(directory, f'Ppub_e_{i}')
        de_file = os.path.join(directory, f'de_{i}')

        sm9_dec(Ppub_e_file, Ke_file, de_file, id_file)
        id_content = read_binary_file(id_file)
        Ke_content = read_binary_file(Ke_file)
        Ppub_e_content = read_binary_file(Ppub_e_file)
        de_content = read_binary_file(de_file)

        # 添加到合并内容
        merged_content += f"用户标识= {id_content}\n"
        merged_content += f"主私钥= {Ke_content}\n"
        merged_content += f"公钥= {Ppub_e_content}\n"
        merged_content += f"私钥= {de_content}\n"
        merged_content += "\n"

    # 写入到txt文件
    if mode == 'ok':
        output_file = 'sm9加密密钥.txt'
    else:
        output_file = 'sm9加密密钥生成速率.txt'


    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(merged_content)
    
    print(f"合并完成，结果已保存到 {output_file}")

if __name__ == "__main__":
    main()
